/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.cassandra.distributed.test.jmx;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeSet;
import javax.management.JMRuntimeException;
import javax.management.MBeanAttributeInfo;
import javax.management.MBeanInfo;
import javax.management.MBeanOperationInfo;
import javax.management.MBeanServerConnection;
import javax.management.ObjectName;
import javax.management.remote.JMXConnector;

import com.google.common.collect.ImmutableSet;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.cassandra.distributed.Cluster;
import org.apache.cassandra.distributed.api.IInstanceConfig;
import org.apache.cassandra.distributed.api.IInvokableInstance;
import org.apache.cassandra.distributed.shared.JMXUtil;

public class JMXTestsUtil
{
    private static final Logger logger = LoggerFactory.getLogger(JMXTestsUtil.class);
    private static final Set<String> IGNORE_ATTRIBUTES = ImmutableSet.of(
    "org.apache.cassandra.net:type=MessagingService:BackPressurePerHost", // throws unsupported saying the feature was removed... dropped in CASSANDRA-15375
    "org.apache.cassandra.db:type=DynamicEndpointSnitch:Scores", // when running in multiple-port-one-IP mode, this fails
    "org.apache.cassandra.db:type=StorageService:OutstandingSchemaVersions", // deprecated (TCM)
    "org.apache.cassandra.db:type=StorageService:OutstandingSchemaVersionsWithPort" // deprecated (TCM)
    );
    private static final Set<String> IGNORE_OPERATIONS = ImmutableSet.of(
    "org.apache.cassandra.db:type=StorageService:stopDaemon", // halts the instance, which then causes the JVM to exit
    "org.apache.cassandra.db:type=StorageService:drain", // don't drain, it stops things which can cause other APIs to be unstable as we are in a stopped state
    "org.apache.cassandra.db:type=StorageService:stopGossiping", // if we stop gossip this can cause other issues, so avoid
    "org.apache.cassandra.db:type=StorageService:resetLocalSchema", // this will fail when there are no other nodes which can serve schema
    "org.apache.cassandra.db:type=StorageService:joinRing", // Causes bootstrapping errors
    "org.apache.cassandra.db:type=Tables,keyspace=system,table=local:loadNewSSTables", // Shouldn't attempt to load SSTables as sometimes the temp directories don't work
    "org.apache.cassandra.db:type=CIDRGroupsMappingManager:loadCidrGroupsCache", // CIDR cache isn't enabled by default
    "org.apache.cassandra.db:type=StorageService:clearConnectionHistory", // Throws a NullPointerException
    "org.apache.cassandra.db:type=StorageService:startGossiping", // causes multiple loops to fail
    "org.apache.cassandra.db:type=StorageService:startNativeTransport", // causes multiple loops to fail
    "org.apache.cassandra.db:type=StorageService:resumeMove", // throws since there is no move in progress
    "org.apache.cassandra.db:type=StorageService:abortMove", // throws since there is no move in progress
    "org.apache.cassandra.db:type=CIDRGroupsMappingManager:loadCidrGroupsCache", // AllowAllCIDRAuthorizer doesn't support this operation, as feature is disabled by default
    "org.apache.cassandra.db:type=StorageService:forceRemoveCompletion", // deprecated (TCM)
    "org.apache.cassandra.db:type=StorageService:createEpochUnsafe", // for Accord testing, but will likely be removed
    "org.apache.cassandra.metrics:type=Client,name=RequestsSizeByIpDistribution:bucketsId" // for Accord testing, but will likely be removed
    );
    // This set of mbeans are registered early enough during the startup of a
    // Cassandra instance for in-jvm dtests to avoid missing registration of mbeans.
    // We ignore both "org.apache.cassandra.diag:type=DiagnosticEventService" and
    // "org.apache.cassandra.diag:type=LastEventIdBroadcaster" because they are being intialized
    // outside the scope of the in-jvm Instance initialization.
    private static final Set<String> EXPECTED_MBEANS_TO_BE_REGISTERED = Set.of(
    "org.apache.cassandra.db:type=EndpointSnitchInfo",
    "org.apache.cassandra.db:type=LocationInfo"
    );

    /**
     * Tests JMX getters and operations and allows passing JMX Env used for the client JMX connection.
     *
     * @param cluster the cluster to test
     * @throws Exception several kinds of exceptions can be thrown, mostly from JMX infrastructure issues.
     */
    public static void testAllValidGetters(Cluster cluster, Map<String, ?> jmxEnv) throws Exception
    {
        Set<String> missingExpectedMbeans = new HashSet<>(EXPECTED_MBEANS_TO_BE_REGISTERED);
        for (IInvokableInstance instance : cluster)
        {
            if (instance.isShutdown())
            {
                continue;
            }
            IInstanceConfig config = instance.config();
            List<Named> errors = new ArrayList<>();
            try (JMXConnector jmxc = JMXUtil.getJmxConnector(config, jmxEnv))
            {
                MBeanServerConnection mbsc = jmxc.getMBeanServerConnection();
                Set<ObjectName> metricNames = new TreeSet<>(mbsc.queryNames(null, null));
                for (ObjectName name : metricNames)
                {
                    if (!name.getDomain().startsWith("org.apache.cassandra"))
                        continue;
                    missingExpectedMbeans.remove(name.getCanonicalName());
                    MBeanInfo info = mbsc.getMBeanInfo(name);
                    for (MBeanAttributeInfo a : info.getAttributes())
                    {
                        String fqn = String.format("%s:%s", name, a.getName());
                        if (!a.isReadable() || IGNORE_ATTRIBUTES.contains(fqn))
                            continue;
                        try
                        {
                            mbsc.getAttribute(name, a.getName());
                        }
                        catch (JMRuntimeException e)
                        {
                            errors.add(new Named(String.format("Attribute %s", fqn), e.getCause()));
                        }
                    }

                    for (MBeanOperationInfo o : info.getOperations())
                    {
                        String fqn = String.format("%s:%s", name, o.getName());
                        if (o.getSignature().length != 0 || IGNORE_OPERATIONS.contains(fqn))
                            continue;
                        try
                        {
                            mbsc.invoke(name, o.getName(), new Object[0], new String[0]);
                        }
                        catch (JMRuntimeException e)
                        {
                            errors.add(new Named(String.format("Operation %s", fqn), e.getCause()));
                        }
                    }
                }
            }
            if (!errors.isEmpty() || !missingExpectedMbeans.isEmpty())
            {
                AssertionError root = new AssertionError();
                for (Named error : errors)
                {
                    // The Named object's message has the cause also so this only logs the message
                    logger.error("Error {}", error.getMessage());
                    root.addSuppressed(error);
                }
                for (String missingMbean : missingExpectedMbeans)
                {
                    // The Named object's message has the cause also so this only logs the message
                    String errorMessage = String.format("Expected mbean %s was not found", missingMbean);
                    logger.error(errorMessage);
                    root.addSuppressed(new RuntimeException(errorMessage));
                }
                throw root;
            }
        }
    }

    /**
     * This class is meant to make new errors easier to read, by adding the JMX endpoint, and cleaning up the unneeded JMX/Reflection logic cluttering the stacktrace
     */
    static class Named extends RuntimeException
    {
        public Named(String msg, Throwable cause)
        {
            super(msg + "\nCaused by: " + cause.getClass().getCanonicalName() + ": " + cause.getMessage(), cause.getCause());
            StackTraceElement[] stack = cause.getStackTrace();
            List<StackTraceElement> copy = new ArrayList<>();
            for (StackTraceElement s : stack)
            {
                if (!s.getClassName().startsWith("org.apache.cassandra"))
                    break;
                copy.add(s);
            }
            Collections.reverse(copy);
            setStackTrace(copy.toArray(new StackTraceElement[0]));
        }
    }
}
